from xtquant import xtdata
import re

def get_financial_futures_code_from_index(index_code:str) -> list:
    """
    传入指数代码，返回对应的期货合约（当前）
    Args:
        index_code:指数代码，如"000300.SH","000905.SH"
    Retuen:
        list: 对应期货合约列表
    """
    # 获取中金所合约
    financial_futures = xtdata.get_stock_list_in_sector("中金所")
    
    # 使用正则表达式筛选合约
    future_list = []
    pattern = r'^[a-zA-Z]{1,2}\d{3,4}\.[A-Z]{2}$'
    for i in financial_futures:
        if re.match(pattern, i):
            future_list.append(i)
    
    # 检查每个合约的底层指数
    ls = []
    for i in future_list:
        _info = xtdata._get_instrument_detail(i)
        
        # 检查合约名称是否包含指数代码信息
        if "InstrumentName" in _info:
            instrument_name = _info["InstrumentName"]
            # 沪深300期货合约名称通常包含"沪深300"
            if index_code == "000300.SH" and "沪深300" in instrument_name:
                ls.append(i)
            # 中证500期货合约名称通常包含"中证500"
            elif index_code == "000905.SH" and "中证500" in instrument_name:
                ls.append(i)
        
        # 检查合约的ProductName是否包含指数信息
        if "ProductName" in _info:
            product_name = _info["ProductName"]
            if index_code == "000300.SH" and ("沪深300" in product_name or "IF" in product_name):
                ls.append(i)
            elif index_code == "000905.SH" and ("中证500" in product_name or "IC" in product_name):
                ls.append(i)
        
        # 检查合约代码前缀
        if index_code == "000300.SH" and i.startswith("IF"):
            ls.append(i)
        elif index_code == "000905.SH" and i.startswith("IC"):
            ls.append(i)
    
    # 去重
    return list(set(ls))

if __name__ == "__main__":
    ls = get_financial_futures_code_from_index("000300.SH")
    print(f"沪深300对应的期货合约: {ls}")
    
    ls2 = get_financial_futures_code_from_index("000905.SH")
    print(f"中证500对应的期货合约: {ls2}")